
import json
import sys

import numpy as np
from scipy.signal import butter, lfilter, find_peaks


def filter(low, hi, data, order=1, yCenter=550, peak_threshold=0, valley_threshold=0, prominence_threshold=0.0,
           min_distance=0):
    f_low = low / 60
    f_hi = hi / 60

    low_wn = f_low * 2 / 50
    hi_wn = f_hi * 2 / 50

    b, a = butter(order, [low_wn, hi_wn], btype='bandpass', analog=False)
    y = lfilter(b, a, data) + yCenter

    peak_valley_info = get_peak_valley(y, peak_threshold, valley_threshold, prominence_threshold, min_distance)

    # 根据原始信号获取频率信息
    freq_info = get_freq(np.array(data))

    ans = np.clip(y, 0, 1100)
    data_arr = np.around(ans)
    return {
        "data": data_arr.tolist(),
        "peakIndex": peak_valley_info["peakIndex"],
        "peakVals": peak_valley_info["peakVals"],
        "valleyIndex": peak_valley_info["valleyIndex"],
        "valleyVals": peak_valley_info["valleyVals"],
        "bpm": freq_info["bpm"]
    }


def get_peak_valley(np_arr, peak_threshold, valley_threshold, prominence_threshold, min_distance):
    peaks, _ = find_peaks(np_arr, height=peak_threshold, prominence=prominence_threshold, distance=min_distance)
    # 找到谷值（将数据取负）
    valleys, _ = find_peaks(-np_arr, height=-valley_threshold, prominence=prominence_threshold,
                            distance=min_distance)
    peaksVal = np.around([np_arr[i] for i in peaks])
    valleysVal = np.around([np_arr[i] for i in valleys])
    return {
        "peakIndex": peaks.tolist(),
        "peakVals": peaksVal.tolist(),
        "valleyIndex": valleys.tolist(),
        "valleyVals": valleysVal.tolist()
    }


def calc_breath_info(peaksVal, peaksIndex):
    # 计算平均值
    if len(peaksVal) <= 1:
        return 0, 0, 0

    last_freq = get_recent_freq_bpm(peaksIndex)

    peak_error = get_error(peaksVal)

    diff = np.diff(peaksIndex)

    freq_diff_arr = 1 / (diff * 0.02)
    freq_error = get_error(freq_diff_arr)

    return round(peak_error, 3), round(freq_error, 3), round(last_freq, 1)


def get_error(arr):
    if len(arr) <= 6:
        return -1
    else:
        last_6_arr = arr[-6:]
        std = np.std(last_6_arr)
        mean = np.mean(last_6_arr)
        cv = std / mean
        return cv


def get_recent_freq_bpm(peaksIndex):
    last_two = peaksIndex[-2:]
    return 1 / (abs(last_two[0] - last_two[1]) * 0.02) * 60


def get_freq(np_arr):
    fft_result = np.fft.fft(np_arr)
    fft_freq = np.fft.fftfreq(len(np_arr))
    # 找到主频率
    main_freq_index = np.argmax(np.abs(fft_result[1:])) + 1
    main_freq = np.abs(fft_freq[main_freq_index])
    each_min_points_count = 60 / 0.02  # 60秒。 每0.02秒产生一个点。
    return {
        "bpm": round(main_freq * each_min_points_count)
    }


if __name__ == "__main__":
    print("setup")  # Signal that the script is ready
    sys.stdout.flush()

    while True:
        try:
            line = sys.stdin.readline().strip()
            if line == "":
                continue
            data = json.loads(line)

            low = data["low"]
            hi = data["hi"]
            content = data["data"]
            order = data['order']
            yCenter = data['yCenter']
            filterStartIndex = data['filterStartIndex']
            peak_threshold = data['peak_threshold']
            valley_threshold = data['valley_threshold']
            prominence_threshold = data['prominence_threshold']
            min_distance = data['min_distance']

            result = filter(low, hi, content, order, yCenter, peak_threshold, valley_threshold, prominence_threshold,
                            min_distance)

            peak_error, freq_error, last_freq = calc_breath_info(result["peakVals"], result["peakIndex"])
            result["filterStartIndex"] = filterStartIndex
            result['peak_error'] = peak_error
            result['freq_error'] = freq_error
            result['last_freq'] = last_freq
            print(json.dumps(result))
            sys.stdout.flush()
        except Exception as e:
            print(json.dumps({"error": str(e)}))
            sys.stdout.flush()
